using Flux


function rank(x::AbstractArray)
    length(size(x))
end


function unsqueeze(x, dims)
    for dim in sort(dims)
        x = Flux.unsqueeze(x, dim)
    end
    x
end


"""
Outer product along the first and the second dimensions of `x` and `y`, respectively.
"""
function outer(x, y)
    unsqueeze(x, [1]) .* unsqueeze(x, [2])

end


function squeeze(x)
    newsize = filter(n -> n > 1, size(x))
    reshape(x, newsize...)
end


function sum_(x; dims)
    sum(x; dims) |> squeeze
end


"""
Assume that the last dimension is the data batch dimension.
"""
function expect(x::AbstractArray{T}) where T<:Real
    last_dim = rank(x)
    Flux.mean(x, dims=last_dim) |> squeeze
end


"""
Assume that the last dimension is the data batch dimension.
"""
function covariance(x::AbstractArray{T}) where T<:Real
    last_dim = rank(x)
    cov(x; dims=last_dim, corrected=false)
end


"""
Inverse function of sigmoid.
"""
function invσ(x)
    log(x) - log(1-x)
end


function hardσ(x)
    if x > 0
        one(x)
    else
        zero(x)
    end
end


function flatten(x)
    reshape(x, prod(size(x)))
end


function L∞(x)
    max(abs.(x)...)
end


# TODO: T<:Float ?
"""
Auxillary function of `zoomin`.
"""
function sample_bernoulli(p::T) where T<:Real
    if rand(typeof(p)) < p
        one(p)
    else
        zero(p)
    end
end


"""
Pearson correlation coefficients.

Parameters
----------
TODO


Returns
-------
Matrix with shape (size(X, 1), size(Y, 1))
"""
function pcc(X::AbstractMatrix{T}, Y::AbstractMatrix{T}) where T<:Real
    σx = var(X; dims=2, corrected=false) .^ 0.5 .+ eps()
    σy = var(Y; dims=2, corrected=false) .^ 0.5 .+ eps()
    C = cov(X, Y; dims=2, corrected=false)

    ρ = zero.(C)
    for i = 1:size(C, 1)
        for j = 1:size(C, 2)
            ρ[i, j] = C[i, j] / σx[i] / σy[j]
        end
    end
    ρ
end


function pcc(X::AbstractMatrix{T}) where T<:Real
    pcc(X, X)
end


function vanish_diag(m::AbstractMatrix)
    m = copy(m)
    for i = 1:size(m, 1)
        for j = 1:size(m, 2)
            if i == j
                m[i, j] = 0
            end
        end
    end
    m
end


function collect_nondiag(m::AbstractMatrix)
    nondiag = eltype(m)[]
    for i = 1:size(m, 1)
        for j = 1:size(m, 2)
            if i != j
                nondiag = push!(nondiag, m[i, j])
            end
        end
    end
    nondiag
end


function ensure_directory(path::String)
    parts = splitdir(path)
    subpath = parts[1]
    for i = 2:size(parts, 1)
        subpath = "$subpath/$(parts[i])"
        try
            mkdir(subpath)
        catch  # has been existed.
            nothing
        end
    end
end